from elasticsearch import helpers

from backend.app.database.db_es import es_conn
from backend.app.schemas.device_v2 import DeviceDetailReq, DeviceSchemaV2, \
    DeviceListReq, DeviceDeleteReq, DeviceBatchDeleteReq, DeleteShopReq, DeviceUpdateOrCreateReq, DeviceBatchOperateReq
from backend.app.utils.const import get_index_by_platform


class GoodsVectorService(object):
    """设备向量化服务"""
    client = None
    embedding = None

    def __init__(self):
        #  获取连接
        if self.client is None:
            self.client = es_conn.get_client()
        if not self.embedding:
            self.embedding = es_conn.get_embeddings()

    def check_exists(self, index_name: str, is_created: bool = True):
        client = self.client
        is_exist = client.indices.exists(index=index_name)
        if not is_exist and is_created:
            client.indices.create(index=index_name, body=mapping)
            is_exist = True
        return is_exist

    @staticmethod
    def get_index_by_platform(platform: str):
        """根据用户信息获取索引名称"""
        index_name = get_index_by_platform(platform)
        return index_name

    def format_vec_req_body(self, req: DeviceDetailReq):
        query_vector: [float] = []
        # 此处判断是否有向量，活着转化为向量
        if req.device_name:
            query_vector = self.embedding.query_vector(req.device_name)
        query_body = {
            "size": 5,
            "query": {
                "script_score": {
                    "query": {
                        "match_all": {}
                    },
                    "script": {
                        "source": "cosineSimilarity(params.query_vector, 'vector') + 1.0",
                        "params": {
                            "query_vector": query_vector
                        }
                    }
                }
            }
        }
        return query_body

    @staticmethod
    def format_search_req(req: DeviceListReq):
        """组装ES请求参数"""
        query_body = {
            "query": {
                "bool": {}
            }
        }
        must_list = []
        if req.device_name:
            must_list.append({"match": {"metadata.device_name": req.device_name}})
        if req.gid:
            must_list.append({"term": {"metadata.gid": str(req.gid)}})

        if must_list:
            query_body["query"]["bool"] = {"filter": must_list}
        else:
            query_body["query"] = {"match_all": {}}

        return query_body

    async def search_list(self, req: DeviceListReq) -> [DeviceSchemaV2]:
        index_name = self.get_index_by_platform(req.platform)
        # 检测索引是否存在
        self.check_exists(index_name)
        req_body = self.format_search_req(req)
        client = self.client

        start = (req.page - 1) * req.size
        if start < 0:
            start = 0

        res = client.search(index=index_name, body=req_body, from_=start, size=req.size)
        if not res.body or res.body["hits"]["total"] == 0:
            return
        device_list: [DeviceSchemaV2] = []
        for device in res.body['hits']['hits']:
            item = DeviceSchemaV2(**device["_source"]["metadata"])
            device_list.append(item)

        return device_list

    async def get_detail_by_id(self, req: DeviceDetailReq) -> DeviceSchemaV2 | None:
        index_name = self.get_index_by_platform(req.platform)
        # 检测索引是否存在
        self.check_exists(index_name)
        client = self.client
        res = client.get(index=index_name, id=req.gid)
        if not (res.body and res.body["found"]):
            return
        obj = DeviceSchemaV2(**res.body["_source"]["metadata"])
        return obj

    def format_to_vec(self, desc: str) -> [float]:
        """将文本向量化"""
        return self.embedding.embed_query(desc)

    async def create_or_update(self, req: DeviceUpdateOrCreateReq):
        client = self.client
        index_name = self.get_index_by_platform(req.platform)
        self.check_exists(index_name)
        #  向量化
        text = req.to_vector_desc()
        vector = self.format_to_vec(text)
        # 将数据转换为字典
        metadata = req.model_dump()
        body = dict(text=text, embedding=vector, metadata=metadata)

        resp = dict(success=True)
        try:
            res = client.index(index=index_name, body=body, id=str(req.gid))
        except Exception as e:
            resp["success"] = False
            resp["msg"] = "更新失败," + str(e)
            return resp
        resp["msg"] = res.body["result"]
        if res.body and res.body["result"]:
            resp["success"] = res.body["_shards"]["successful"] == 1
        return resp

    async def batch_create_or_update(self, req: DeviceBatchOperateReq):
        client = self.client
        index_name = self.get_index_by_platform(req.platform)
        goods_list = []
        for good in req.device_list:
            #  向量化
            text = good.to_vector_desc()
            vector = self.format_to_vec(text)
            metadata = good.model_dump()
            body = dict(text=text, embedding=vector, metadata=metadata)
            item = {
                "_op_type": "update",
                "_index": index_name,
                "_id": good.gid,
                "doc": body,
                "doc_as_upsert": True  # 如果文档不存在，插入它
            }
            goods_list.append(item)

        resp = dict(success=True)
        try:
            res = helpers.bulk(client, goods_list)
        except Exception as e:
            resp["success"] = False
            resp["msg"] = "更新失败," + str(e)
            return resp
        resp["msg"] = "更新成功"
        resp["data"] = res
        return resp

    async def delete_by_id(self, req: DeviceDeleteReq):
        index_name = self.get_index_by_platform(req.platform)
        client = self.client
        resp = dict(success=True)
        try:
            res = client.delete(index=index_name, id=req.gid)
        except Exception as e:
            resp["success"] = False
            resp["msg"] = "删除失败，" + str(e)
            return resp

        resp["msg"] = res.body["result"]
        if res.body and res.body["result"]:
            resp["success"] = res.body["_shards"]["successful"] == 1
        return resp

    async def batch_delete(self, req: DeviceBatchDeleteReq):
        index_name = self.get_index_by_platform(req.platform)
        client = self.client
        resp = dict(success=True)
        delete_req = []
        for gid in req.gids:
            item = {"_op_type": "delete", "_index": index_name, "_id": gid}
            delete_req.append(item)
        if not delete_req:
            return

        try:
            res = helpers.bulk(client, delete_req)
        except Exception as e:
            resp["success"] = False
            resp["msg"] = "删除失败，" + str(e)
            return resp
        resp["msg"] = "删除成功"
        return resp

    def delete_shop(self, req: DeleteShopReq):
        index_name = self.get_index_by_platform(req.platform)
        client = self.client
        resp = dict(success=True)
        if not self.check_exists(index_name, is_created=False):
            return resp
        try:
            res = client.indices.delete(index=index_name)
        except Exception as e:
            resp["success"] = False
            resp["msg"] = "删除失败，" + str(e)
            return resp
        resp["data"] = res
        resp["msg"] = "删除成功"
        return resp


# 索引数据结构
device_mapping = {
    "mappings": {
        "properties": {
            "vector": {
                "type": "dense_vector",
                "dims": 1536
            },
            "gid": {
                "type": "keyword"
            },
            "sid": {
                "type": "keyword"
            },
            "device_name": {
                "type": "text"
            },
            "device_avatar": {
                "type": "text",
                "index": False
            },
            "mod1": {
                "type": "text"
            },
            "mod": {
                "type": "keyword"
            },
            "brand_id": {
                "type": "keyword"
            },
            "brand_title": {
                "type": "text"
            },
            "class_id": {
                "type": "keyword"
            },
            "class_name": {
                "type": "text"
            },
            "class_1_id": {
                "type": "keyword"
            },
            "class_1_name": {
                "type": "text"
            },
            "class_2_id": {
                "type": "keyword"
            },
            "class_2_name": {
                "type": "text"
            },
            "price": {
                "type": "text",
                "index": False
            },
            "meter": {
                "type": "text",
            },
            "meter_clean": {
                "type": "text",
                "index": False
            },
            "spare_meter": {
                "type": "text",
                "index": False
            },
            "address": {
                "type": "text"
            },
            "meter_coding": {
                "type": "keyword"
            },
            "country_title": {
                "type": "text"
            },
            "country": {
                "type": "text"
            },
            "text": {
                "type": "text"
            },
            "metadata": {
                "type": "text"
            }
        }
    }
}

mapping = {
    "mappings": {
        "properties": {
            "text": {"type": "text"},
            "embedding": {"type": "dense_vector", "dims": 1536, "index": True, "similarity": "cosine"},  # dense_vector 字段，假设维度是 128
            "metadata": {
                "type": "object",
                "properties": {
                    "gid": {"type": "keyword"},
                    "sid": {"type": "keyword"},
                    "device_name": {"type": "text"},
                    "device_avatar": {"type": "text"},
                    "mod1": {"type": "text"},
                    "mod": {"type": "keyword"},
                    "brand_id": {"type": "keyword"},
                    "brand_title": {"type": "text"},
                    "class_id": {"type": "keyword"},
                    "class_name": {"type": "text"},
                    "class_1_id": {"type": "keyword"},
                    "class_1_name": {"type": "text"},
                    "class_2_id": {"type": "keyword"},
                    "class_2_name": {"type": "text"},
                    "price": {"type": "text"},
                    "meter": {"type": "text"},
                    "meter_clean": {"type": "text"},
                    "spare_meter": {"type": "text"},
                    "address": {"type": "text"},
                    "meter_coding": {"type": "keyword"},
                    "country_title": {"type": "text"},
                    "country": {"type": "text"}
                }
            }
        }
    }
}
